# coding = utf-8

'''
针对nnUnet的一些分析代码
'''

import SimpleITK as sitk
import pickle
import numpy as np
import matplotlib.pyplot as plt

'''
分析预处理之后的数据
'''
def analysis_preprocessing_data():
    raw_data = "/datasets/nnUNet/raw_data_base/nnUNet_raw_data/Task101_Kidney/imagesTr/case_00000_0000.nii.gz"
    raw_label = "/datasets/nnUNet/raw_data_base/nnUNet_raw_data/Task101_Kidney/labelsTr/case_00000.nii.gz"

    raw_data = sitk.ReadImage(raw_data)
    raw_label = sitk.ReadImage(raw_label)

    raw_image_array = sitk.GetArrayFromImage(raw_data)
    raw_label_array = sitk.GetArrayFromImage(raw_label)

    raw_image_array = raw_image_array.transpose((2,0,1))
    raw_label_array = raw_label_array.transpose((2,0,1))

    stage0_data_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/nnUNetData_plans_v2.1_stage0/case_00001.npz"
    stage1_data_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/nnUNetData_plans_v2.1_stage1/case_00001.npz"

    stage0_data = np.load(stage0_data_file)["data"]
    stage1_data = np.load(stage1_data_file)["data"]

    stage0_image = stage0_data[0]
    stage0_label = stage0_data[1]

    stage1_image = stage1_data[0]
    stage1_label = stage1_data[1]

    print(np.unique(raw_label_array))
    print(np.unique(stage0_label))
    print(np.unique(stage1_label))

    for i in range(stage1_label.shape[0]):
        if -1 in np.unique(stage1_label[i]):
            print(i)

    print(stage1_label[0])



'''
读取pkl文件
'''
def analysisPKL():
    dataset_properties_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/dataset_properties.pkl"
    model_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/nnUNetPlansv2.1_plans_3D.pkl"

    dataset_properties = pickle.load(open(dataset_properties_file, "rb"))
    model = pickle.load(open(model_file, "rb"))

    stage0_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/nnUNetData_plans_v2.1_stage0/case_00001.pkl"
    stage1_file = "/datasets/nnUNet/preprocessed/Task101_Kidney/nnUNetData_plans_v2.1_stage1/case_00001.pkl"

    stage0 = pickle.load(open(stage0_file,"rb"))
    stage1 = pickle.load(open(stage1_file, "rb"))

    for key in stage0.keys():
        print(key, stage0[key], stage1[key])




if __name__ == '__main__':
    analysis_preprocessing_data()